#!/usr/bin/env python3
"""检查 ChromaDB 集合和数据"""

import chromadb
import os
from pathlib import Path

def check_chromadb_data():
    """检查 ChromaDB 数据"""
    print("=== 检查 ChromaDB 集合和数据 ===")
    
    # 检查两个可能的数据目录
    data_dirs = [
        "./data/chroma_db",
        "./data/chromadb"
    ]
    
    for data_dir in data_dirs:
        print(f"\n检查数据目录: {data_dir}")
        if not os.path.exists(data_dir):
            print(f"❌ 目录不存在: {data_dir}")
            continue
            
        try:
            # 连接到 ChromaDB
            client = chromadb.PersistentClient(path=data_dir)
            
            # 获取所有集合
            collections = client.list_collections()
            print(f"✅ 找到 {len(collections)} 个集合:")
            
            for collection in collections:
                print(f"  - 集合名: {collection.name}")
                print(f"    ID: {collection.id}")
                
                # 获取集合中的文档数量
                try:
                    count = collection.count()
                    print(f"    文档数量: {count}")
                    
                    if count > 0:
                        # 获取一些示例文档
                        results = collection.get(limit=3)
                        print(f"    示例文档:")
                        for i, (doc_id, document, metadata) in enumerate(zip(
                            results['ids'], 
                            results['documents'], 
                            results['metadatas']
                        )):
                            print(f"      [{i+1}] ID: {doc_id}")
                            print(f"          内容: {document[:100]}..." if len(document) > 100 else f"          内容: {document}")
                            print(f"          元数据: {metadata}")
                            
                        # 如果是 conversations 集合，检查用户ID
                        if collection.name == "conversations":
                            print(f"    检查用户ID分布:")
                            all_results = collection.get()
                            user_ids = set()
                            for metadata in all_results['metadatas']:
                                if metadata and 'user_id' in metadata:
                                    user_ids.add(metadata['user_id'])
                            print(f"    用户ID: {list(user_ids)}")
                            
                            # 测试搜索
                            print(f"    测试搜索 '王者荣耀':")
                            search_results = collection.query(
                                query_texts=["王者荣耀"],
                                n_results=3
                            )
                            print(f"    搜索结果数量: {len(search_results['ids'][0])}")
                            
                            # 测试带用户过滤的搜索
                            print(f"    测试搜索 '王者荣耀' (用户: default_user):")
                            search_results_filtered = collection.query(
                                query_texts=["王者荣耀"],
                                n_results=3,
                                where={"user_id": "default_user"}
                            )
                            print(f"    过滤搜索结果数量: {len(search_results_filtered['ids'][0])}")
                            
                except Exception as e:
                    print(f"    ❌ 获取集合信息失败: {e}")
                    
        except Exception as e:
            print(f"❌ 连接 ChromaDB 失败: {e}")

if __name__ == "__main__":
    check_chromadb_data()